 /**
 * @file kruskal.cpp
 * @author Wang Bo (1276106221@qq.com)
 * @brief
 * Kruskal 最小生成树问题（MST）
 * 贪心算法典型应用
 * 并查集数据结构典型应用
 * @version 0.1
 * @date 2021-07-07
 * @copyright Copyright (c) 2021
 */

#include <iostream>
#include <vector>
#include <queue>
#include <map>
#include <algorithm>
#include <random>
#include <numeric>
#include <functional>

using std::cout;
using std::endl;
using std::vector;

//并查集数据结构
template <typename T>
class DisJointSet {
public:
    DisJointSet(size_t N) {
        nodes.reserve(N);
    }
    //向并查集中输入顶点
    void addSet(const T& data) {
        nodes.emplace_back(data);
        parant.emplace_back(nodes.size() - 1); //新加入的点的根节点都是自己
        rank.emplace_back(0); //并且秩(也就是子节点是数目)为0
    }
    //查询数据的根节点
    size_t find(const T& data) {
        auto it = std::find_if(nodes.begin(), nodes.end(), [&](const Node& n){
            return n.data_ == data;
        });
        if (it == nodes.end()) //异常处理
            return size_t(-1);
        auto nodeIdx = std::distance(nodes.begin(), it);
        auto parentIdx = parant[nodeIdx];
        while (parentIdx != nodeIdx) {
            nodeIdx = parentIdx;
            parentIdx = parant[nodeIdx];
        }
        return parentIdx;
    }
    //合并x和y两个集合
    void unionSet(const T& x, const T& y)
    {
        auto fx = find(x);
        auto fy = find(y);
        if (fx == fy) {
            return;
        }
        else if (rank[fx] > rank[fy]) {
            //如果x所在的树比y所在的树高
            parant[fy] = parant[fx]; //则将y树根节点连接到x树根节点上
            rank[fx] += 1; //x树根节点的秩+1
        }
        else {
            parant[fx] = parant[fy]; //否则将x树连接到y树上
            rank[fy] += 1;
        }
    }

private:
    struct Node {
        T data_;
        Node(const T& data): data_(data) {}
        inline bool operator!=(const Node& that) const {
            return data_ != that.data_;
        }
    };
    vector<Node> nodes;
    vector<size_t> parant;
    vector<size_t> rank;
};

template <typename T>
struct Edge { //图的边结构
    size_t src; //源节点下标
    size_t dest; //目地节点下标
    T weight; //权重
    inline bool operator<(const Edge<T>& that) const {
        return weight < that.weight;
    }
    inline bool operator>(const Edge<T>& that) const {
        return weight > that.weight;
    }
};

//图结构，注意顶点是从1~N的（不是0~N-1）
template <typename T>
class Graph {
public:
    Graph(size_t N): vertices_(N) {}
    //返回图的顶点数
    auto vertices() const { return vertices_; }
    //返回只读的边列表
    const auto& edges() const { return edgeList_; }
    void addEdge(const Edge<T>& e) {
        if (e.src >= 1 && e.src <= vertices_ && e.dest >= 1 && e.dest <= vertices_)
            edgeList_.emplace_back(e);
        else
            std::cerr << "Vertex out of range\n";
    }
    void addEdge(Edge<T>&& e) {
        if (e.src >= 1 && e.src <= vertices_ && e.dest >= 1 && e.dest <= vertices_)
            edgeList_.emplace_back(std::move(e));
        else
            std::cerr << "Vertex out of range\n";
    }
    //获取从顶点v出发可以走的边
    auto edgesFromVertex(size_t v) const {
        vector<Edge<T>> edgesRet;
        for (auto& e : edgeList_) {
            if (e.src == v)
                edgesRet.emplace_back(e);
        }
        return edgesRet;
    }
    //友元输出函数，依次输出个顶点，与其每条边所连接的终点及边的权重
    friend std::ostream& operator<<(std::ostream& os, const Graph<T>& g) {
        for (auto i = 1; i <= g.vertices(); ++i) {
            auto edgesFromVi = g.edgesFromVertex(i);
            os << i << ":\t";
            for (auto& e : edgesFromVi) {
                os << "{" << e.dest << ": " << e.weight << "}, ";
            }
            os << endl;
        }
        return os;
    }
private:
    size_t vertices_; //顶点数量
    vector<Edge<T>> edgeList_;
};

//最小生成树是一个树形结构，因此也可以复用图Graph的结构来表示
template <typename T>
Graph<T> MSTSolution(const Graph<T>& G) {
    //首先建立一个边的优先队列（小根堆），按照边的权重建堆
    std::priority_queue<Edge<T>, vector<Edge<T>>, std::greater<Edge<T>>> minEdgeHeap;
    for (const auto& e : G.edges()) {
        minEdgeHeap.push(e);
    }
    auto N = G.vertices();
    DisJointSet<size_t> djs(N);
    //第一步：将所有顶点1~N加入并查集
    for (size_t i = 1; i <= N; ++i) {
        djs.addSet(i);
    }
    Graph<T> MST(N);
    //不断从小根堆中取出边
    while (!minEdgeHeap.empty()) {
        auto e = std::move(minEdgeHeap.top());
        minEdgeHeap.pop();
        //判断边所连接的两个顶点是否已经在并查集中建立连接了
        if (djs.find(e.src) != djs.find(e.dest)) {
            //若尚未连接，将这个边的两个顶点在并查集中合并(连接)
            djs.unionSet(e.src, e.dest);
            //在最小生成树中存入这条边
            MST.addEdge(std::move(e));
        }
        //若已连接，则丢弃这条边
    }
    return MST;
}

int main(int argc, const char** argv) {
    using T = unsigned;

	Graph<T> G(8);

	std::map<unsigned, std::vector<std::pair<size_t, T>>> edges;
	edges[1] = { {2, 2}, {5, 3} };
	edges[2] = { {1, 2}, {5, 5}, {4, 1} };
	edges[3] = { {4, 2}, {7, 3} };
	edges[4] = { {2, 1}, {3, 2}, {5, 2}, {6, 4}, {8, 5} };
	edges[5] = { {1, 3}, {2, 5}, {4, 2}, {8, 3} };
	edges[6] = { {4, 4}, {7, 4}, {8, 1} };
	edges[7] = { {3, 3}, {6, 4} };
	edges[8] = { {4, 5}, {5, 3}, {6, 1} };

	for (auto& i : edges)
		for(auto& j: i.second)
			G.addEdge(Edge<T>{ i.first, j.first, j.second });

	std::cout << "Original Graph" << std::endl;
	std::cout << G;

	auto MST = MSTSolution(G);
	std::cout << std::endl << "Minimum Spanning Tree" << std::endl;
	std::cout << MST;
    return 0;
}